#include "RSAUtil.h"

int rsa_public_encrypt(const char *public_key_path,
                           unsigned char *in, size_t in_len,
                           unsigned char **out, size_t *out_len){
    int ret = 0;
    size_t encrypt_buf_len = 0;
    unsigned char *encrypt_buf = NULL;
    
    EVP_PKEY_CTX *ctx = NULL;
    ENGINE *eng = NULL;
    int pad = RSA_PKCS1_PADDING;
    EVP_PKEY *pub_key = NULL;
    BIO *key = NULL;
    BIO *bio_err = BIO_new_fp(stderr, BIO_NOCLOSE);
    
    // BIO *BIO_new_file(const char *filename, const char *mode);
    if((key = BIO_new_file(public_key_path, "rb")) == NULL){
        BIO_printf(bio_err, "unable to read key %s\n", public_key_path);
        ret = -1;
        goto cleanup;
    }
    // RSA *PEM_read_bio_RSA_PUBKEY(BIO *bp, RSA **x, pem_password_cb *cb, void *u);
    if ((pub_key = PEM_read_bio_PUBKEY(key, NULL, (pem_password_cb *)NULL, NULL)) == NULL){
        BIO_printf(bio_err, "PEM_read_bio_PUBKEY error\n");
        ret = -2;
        goto cleanup;
    }

    // EVP_PKEY_CTX *EVP_PKEY_CTX_new(EVP_PKEY *pkey, ENGINE *e);
    if ((ctx = EVP_PKEY_CTX_new(pub_key, eng)) == NULL){
        fprintf(stderr, "EVP_PKEY_CTX_new error\n");
        ret = -3;
        goto cleanup;
    }
    if (EVP_PKEY_encrypt_init(ctx) <= 0){
        fprintf(stderr, "EVP_PKEY_encrypt_init error\n");
        ret = -4;
        goto cleanup;
    }
    // int EVP_PKEY_CTX_set_rsa_padding(EVP_PKEY_CTX *ctx, int pad);
    if (EVP_PKEY_CTX_set_rsa_padding(ctx, pad) <= 0){
        fprintf(stderr, "EVP_PKEY_CTX_set_rsa_padding error\n");
        ret = -5;
        goto cleanup;
    }
    
    // int EVP_PKEY_encrypt(EVP_PKEY_CTX *ctx, unsigned char *out, size_t *outlen, const unsigned char *in, size_t inlen);
    if (EVP_PKEY_encrypt(ctx, NULL, &encrypt_buf_len, in, in_len) <= 0){
        fprintf(stderr, "EVP_PKEY_encrypt determine buffer length error\n");
        ret = -6;
        goto cleanup;
    }
    if ((encrypt_buf = OPENSSL_malloc(encrypt_buf_len))  == NULL) {
        fprintf(stderr, "Malloc failed.\n");
        ret = -7;
        goto cleanup;
    }
    if (EVP_PKEY_encrypt(ctx, encrypt_buf, &encrypt_buf_len, in, in_len) <= 0) {
        fprintf(stderr, "EVP_PKEY_encrypt() failed.\n");
        ret = -8;
        goto cleanup;
    }
    
    if(*out == NULL){
        *out = (unsigned char *)malloc(BUFFER_SIZE);
        if (*out  == NULL) {
            fprintf(stderr, "Malloc failed.\n");
            ret = -9;
            goto cleanup;
        }
        memset(*out, '\0', BUFFER_SIZE);
    }
#ifdef BASE64_WITHIN
    // output with base64
    unsigned char *b4_buf = (unsigned char *)malloc(BUFFER_SIZE);
    memset(b4_buf, '\0', BUFFER_SIZE);
    base64_encode(encrypt_buf, (int)encrypt_buf_len, (unsigned char *)*out, (int *)out_len);
    logger_info_fmt("base64:[%s]\n", (*out));
#else
    *out_len = (int)encrypt_buf_len;
    strncpy((char *)(*out), (const char *)encrypt_buf, (int)encrypt_buf_len);
#endif
    
#ifdef ENTRY_TEST
    fprintf(stdout, "Encrypt:\n");
    BIO_dump_indent_fp(stdout, (const char *)(*out), *out_len, 2);
#endif
    
cleanup:
    if (!ret)
        OPENSSL_free(encrypt_buf);
    BIO_free(key);
    EVP_PKEY_free(pub_key);
    EVP_PKEY_CTX_free(ctx);
    return ret;
}

int rsa_private_decrypt(const char *private_key_path,
                           unsigned char *in, size_t in_len,
                           unsigned char **out, size_t *out_len){
    int ret = 0;
    size_t decrypt_buf_len = 0;
    unsigned char *decrypt_buf = NULL;
    
    EVP_PKEY_CTX *ctx = NULL;
    ENGINE *eng = NULL;
    int pad = RSA_PKCS1_PADDING;
    EVP_PKEY *pri_key = NULL;
    BIO *key = NULL;
    BIO *bio_err = BIO_new_fp(stderr, BIO_NOCLOSE);
    
    // BIO *BIO_new_file(const char *filename, const char *mode);
    if((key = BIO_new_file(private_key_path, "rb")) == NULL){
        BIO_printf(bio_err, "unable to read key %s\n", private_key_path);
        ret = -1;
        goto cleanup;
    }
    //  EVP_PKEY *PEM_read_bio_PrivateKey(BIO *bp, EVP_PKEY **x, pem_password_cb *cb, void *u);
    if ((pri_key = PEM_read_bio_PrivateKey(key, NULL, (pem_password_cb *)NULL, NULL)) == NULL){
        BIO_printf(bio_err, "PEM_read_bio_PrivateKey error\n");
        ret = -2;
        goto cleanup;
    }

    // EVP_PKEY_CTX *EVP_PKEY_CTX_new(EVP_PKEY *pkey, ENGINE *e);
    if ((ctx = EVP_PKEY_CTX_new(pri_key, eng)) == NULL){
        fprintf(stderr, "EVP_PKEY_CTX_new error\n");
        ret = -3;
        goto cleanup;
    }
    if (EVP_PKEY_decrypt_init(ctx) <= 0){
        fprintf(stderr, "EVP_PKEY_decrypt_init error\n");
        ret = -4;
        goto cleanup;
    }
    // int EVP_PKEY_CTX_set_rsa_padding(EVP_PKEY_CTX *ctx, int pad);
    if (EVP_PKEY_CTX_set_rsa_padding(ctx, pad) <= 0){
        fprintf(stderr, "EVP_PKEY_CTX_set_rsa_padding error\n");
        ret = -5;
        goto cleanup;
    }
    
    unsigned char *buf;
    int buf_len = 0;
#ifdef BASE64_WITHIN
    // base64 process
    unsigned char *b4_decode_buf = (unsigned char *)malloc(BUFFER_SIZE);
    memset(b4_decode_buf, '\0', BUFFER_SIZE);
    int b4_decode_len = 0;
    base64_decode(in, (int)in_len, (unsigned char *)b4_decode_buf, &b4_decode_len);
    // logger_info_fmt("b4_decode_buf=[%s]\n", b4_decode_buf);
    buf = b4_decode_buf;
    buf_len = b4_decode_len;
#else
    buf = in;
    buf_len = (int)in_len;
#endif
    //  int EVP_PKEY_decrypt(EVP_PKEY_CTX *ctx, unsigned char *out, size_t *outlen, const unsigned char *in, size_t inlen);
    if (EVP_PKEY_decrypt(ctx, NULL, &decrypt_buf_len, buf, buf_len) <= 0){
        fprintf(stderr, "EVP_PKEY_decrypt determine buffer length error\n");
        ret = -6;
        goto cleanup;
    }
    if ((decrypt_buf = OPENSSL_malloc(decrypt_buf_len))  == NULL) {
        fprintf(stderr, "Malloc failed.\n");
        ret = -7;
        goto cleanup;
    }
    if (EVP_PKEY_decrypt(ctx, decrypt_buf, &decrypt_buf_len, buf, buf_len) <= 0) {
        fprintf(stderr, "EVP_PKEY_decrypt() failed.\n");
        ret = -8;
        goto cleanup;
    }
    if(*out == NULL){
        *out = (unsigned char *)malloc(BUFFER_SIZE);
        if (*out  == NULL) {
            fprintf(stderr, "Malloc failed.\n");
            ret = -9;
            goto cleanup;
        }
        memset(*out, '\0', BUFFER_SIZE);
    }
    
    *out_len = strlen((const char *)decrypt_buf);
    strncpy((char *)(*out), (const char *)decrypt_buf, *out_len);
    
#ifdef ENTRY_TEST
    fprintf(stdout, "Decrypt:\n");
    BIO_dump_indent_fp(stdout, (const char *)(*out), *out_len, 2);
#endif
    
cleanup:
    if (!ret)
        OPENSSL_free(decrypt_buf);
    BIO_free(key);
    EVP_PKEY_free(pri_key);
    EVP_PKEY_CTX_free(ctx);
    return ret;
}
 
void printLastError(){
    char * err = (char*)malloc(256);;
    ERR_load_crypto_strings();
    ERR_error_string(ERR_get_error(), err);
    logger_err(err);
    free(err);
}

#ifdef ENTRY_TEST
int main(int argc, char **argv){
    // echo -n '!QAZ2wsx' | base64
    const char *pwd = "!QAZ2wsx"; // IVFBWjJ3c3g=
    unsigned char *encrypted_buffer = NULL;
    unsigned char *decrypted_buffer = NULL;
    int encrypted_size = 0;
    int decrypted_size = 0;
    int ret = 0;


#ifdef BASE64_TEST
    unsigned char *b4_decode_buf = (unsigned char *)malloc(BUFFER_SIZE);
    unsigned char *b4_encode_buf = (unsigned char *)malloc(BUFFER_SIZE);
    memset(b4_decode_buf, '\0', BUFFER_SIZE);
    memset(b4_encode_buf, '\0', BUFFER_SIZE);
    int b4_decode_len = 0;
    int b4_encode_len = 0;
    base64_encode((unsigned char *)pwd, (int)strlen(pwd), (unsigned char *)b4_encode_buf, &b4_encode_len);
    logger_info_fmt("base64_encode: [%s]\n", b4_encode_buf);
    base64_decode((unsigned char *)b4_encode_buf, (int)strlen((const char *)b4_encode_buf), (unsigned char *)b4_decode_buf, &b4_decode_len);
    logger_info_fmt("base64_decode: [%s]\n", b4_decode_buf);
    if(b4_decode_buf) free(b4_decode_buf);
    if(b4_encode_buf) free(b4_encode_buf);
#endif
    
#ifdef BASE64_WITHIN
    logger_info("===========================================================================\n");
    logger_info("Note: RSA base64 inside. \n");
#endif
    logger_info("===========================================================================\n");
    if((ret = rsa_public_encrypt("/Users/jiuzhu/Documents/openssl/src/rsa_public_key_demo.pem", (unsigned char *)pwd, strlen(pwd), &encrypted_buffer, (size_t *)&encrypted_size)) != 0){
        printLastError();
        return -1;
    }

    // echo $(echo -n '!QAZ2wsx' | openssl rsautl -encrypt -pubin -inkey src/rsa_public_key_demo.pem | base64) | base64 -d | openssl rsautl -decrypt -inkey src/rsa_private_key_demo.pem
    // echo $(echo -n $(echo -n '!QAZ2wsx' | base64) | openssl rsautl -encrypt -pubin -inkey src/rsa_public_key_demo.pem | base64)
    // echo $(echo -n $(echo -n '123' | base64 -d) | openssl rsautl -decrypt -inkey rsa_private_key_demo.pem
    logger_info_fmt("==> encrypt ret=%d\n", ret);
    if((ret = rsa_private_decrypt("/Users/jiuzhu/Documents/openssl/src/rsa_private_key_demo.pem", encrypted_buffer, encrypted_size, &decrypted_buffer, (size_t *)&decrypted_size)) != 0){
        printLastError();
        return -2;
    }
    logger_info_fmt("==> decrypt ret=%d\n", ret);
    if(strncmp((const char *)decrypted_buffer, "!QAZ2wsx", strlen("!QAZ2wsx")) != 0){
        logger_err("not match\n");
        exit(-1);
    }
    if(!encrypted_buffer) free(encrypted_buffer);
    if(!decrypted_buffer) free(decrypted_buffer);
    return 0;
}
#endif
